#
# Rabi oscillations of qubit subject to a classical driving field.
#
from qutip import *
from nv_system import *
from pulse_control import *
from rotating_frame import *
from scipy import fftpack    
   
def electron_fid(nvsys,init_state,ob_op,evo_time=100,Bz=0.05,sampling_rate=10,T1=0,T2=0):


    rou0 = init_state     # initial state
    
    #evolution time in ns    
    sampling_times = sampling_rate * evo_time  

    #the time-independent hamiltonian
    H0 = nvsys.get_static_ham(Bz)*pi*2
    Ht=[H0]     
    
    # define the time-dependence of the hamiltonian using the list-string format
    args = {}    

    # unitary evolution
    c_op_list = []
    if T1 and T2:
        c_op_list.extend(nvsys.get_lindblad_op(1, T1, T2))
    
    #options for mesolve
    opts=Odeoptions()
    opts.nsteps=1e6
    
    # evolve and system subject to the time-dependent hamiltonian
    outstates = []

    tlist = linspace(0, evo_time, sampling_times)
    output = mesolve(Ht, rou0, tlist, c_op_list, [ob_op], args,opts)      
    return tlist,output

default_NV_C13_coupling = [
                           [5e-3,-6.3e-3,-2.9e-3],
                           [-6.3e-3,4.2e-3,-2.3e-3],
                           [-2.9e-3,-2.3e-3,8.2e-3]]
default_NV_C13_coupling = [
                           [5e-3,-6.3e-3,-2.9e-3],
                           [-6.3e-3,4.2e-3,-2.3e-3],
                           [-2.9e-3,-2.3e-3,10e-3]]
if __name__=='__main__':
    evo_time = 50 
    nvsys = nv_system([nv_electron_spin(),C13_nuclear_spin()],[[0,1,default_NV_C13_coupling]]) 
    
    state00 = nvsys.get_basis([0,-0.5])
    state01 = nvsys.get_basis([0,0.5])
    state10 = nvsys.get_basis([1,-0.5])
    state11 = nvsys.get_basis([1,0.5])
    B0=[0,0,0.05]
    init_state = tensor(ket2dm(basis(3,0)+basis(3,1)),qeye(2))     # initial state
    rop =tensor( basis(3,0)*basis(3,1).dag() + basis(3,1)*basis(3,0).dag(),qeye(2))
    rot = (-1j*rop * pi/4).expm()
    ob_op = rot * tensor(ket2dm(basis(3,1)),qeye(2)) * rot.dag() 
    #init_state = init_state / init_state.norm()
    tlist,output=electron_fid(nvsys,init_state,ob_op,evo_time = 10000,sampling_rate=10,T1=0,T2=0)
#     transitions = [[basis.states[0],basis.states[1],nvsys.get_transition_freq(basis.states[0],basis.states[1])]]
#     rframe = rotation_frame(transitions)
#     rop = rframe.get_transfer_operator()[0]
#     data = []
#     for i in range(len(tlist)):
#         t = tlist[i]
#         state = rot * states[i] * rot.dag()
#         data.append(expect(tensor(ket2dm(basis(3,1)),qeye(2)),state))
#     
    # Plot the result
    data = output.expect[0]
    subplot(211)
    plot(tlist, real(data))
    xlabel('Time')
    ylabel('Occupation probability with ' +str(1) +" loop")
    title('Excitation probabilty of qubit')
    
    deltat=tlist[1]-tlist[0]
    # remove the DC component
    data = data - average(data) 
    fft= fftpack.fft(data)
    freqs = fftpack.fftfreq(len(data),deltat)
    
    subplot(212)
    plot(freqs,real(fft),'r')
    xlabel('Freq')
    ylabel('Arbitrary Unit')
    title('FID Spectrum')
    
    show()
